#这个文件用来将txt或xml和image生成带圈注的image,YOLOV4的标注txt文件格式
#一个.txt或.xml + 一个.jpg   --------->   一个.jpg
import cv2
import os
from tqdm import tqdm

origin_images_path = "/media/jiji/fe0e60a9-bc54-4761-a52c-dd4ec10ba6db/Company/PSE代码里的数据/syrq/原img/"
origin_label_path = "/media/jiji/fe0e60a9-bc54-4761-a52c-dd4ec10ba6db/Company/PSE代码里的数据/syrq/原txt/"
images_save_path = "/media/jiji/fe0e60a9-bc54-4761-a52c-dd4ec10ba6db/Company/PSE代码里的数据/syrq/我的/"

xml_or_txt = 'txt'

if __name__ == '__main__':
    file_names = os.listdir(origin_label_path)
    if xml_or_txt == 'txt':
        for file_name in tqdm(file_names,ncols=150):

            img = cv2.imread(origin_images_path+file_name[:-4]+".jpg")
            img_h,img_w = img.shape[0],img.shape[1]
            txt_file = open(origin_label_path+file_name,'r')
            for line in txt_file.readlines():
                line_chars = line.split(' ')
                class_id = int(line_chars[0])
                #转换为float
                coordinate_x,coordianete_y = float(line_chars[1]),float(line_chars[2])
                bbox_w,bbox_h = float(line_chars[3]),float(line_chars[4])

                coordinate_x,coordianete_y = round((coordinate_x-bbox_w/2)*img_w),round((coordianete_y-bbox_h/2)*img_h)
                bbox_w,bbox_h = round(bbox_w*img_w),round(bbox_h*img_h)
                point_1,point_2 = [coordinate_x,coordianete_y],[coordinate_x+bbox_w,coordianete_y+bbox_h]
                cv2.rectangle(img,point_1,point_2,(0,0,255),2)
            cv2.imwrite(images_save_path+file_name[:-4]+".jpg",img)

    elif xml_or_txt == 'xml':
        pass
    else:
        raise RuntimeError('参数xml_or_txt错误')
